#include "SrpcChannel.h"
#include "SrpcHeader.pb.h"
#include "SrpcLogger.h"
#include "ZooKeeperUtil.h"
#include <arpa/inet.h>
#include <google/protobuf/descriptor.h>
#include <google/protobuf/service.h>
#include <netinet/in.h>
#include <string>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

std::mutex g_data_mutex;

SrpcChannel::SrpcChannel(bool connectNow) : m_clientfd(-1) {
  if (!connectNow) {
    return;
  }
  // 建立连接
  auto rt = newConnect(m_ip.c_str(), m_port);
  // 多次重试
  int count = 3;
  while (!rt && count--) {
    rt = newConnect(m_ip.c_str(), m_port);
  }
}

SrpcChannel::~SrpcChannel() {}

void SrpcChannel::CallMethod(const google::protobuf::MethodDescriptor *method,
                             google::protobuf::RpcController *controller,
                             const google::protobuf::Message *request,
                             google::protobuf::Message *response,
                             google::protobuf::Closure *done) {
  if (m_clientfd == -1) {
    // 获取服务对象名和方法名
    const google::protobuf::ServiceDescriptor *sd = method->service();
    // 服务名
    m_service_name = sd->name();
    // 方法名
    m_method_name = method->name();
    // 客户端从Zookeeper找到提供该服务的服务器地址
    ZooKeeperClient zkClient;
    zkClient.Start();
    std::string host_data =
        QueryServiceHost(&zkClient, m_service_name, m_method_name, m_idx);
    m_ip = host_data.substr(0, m_idx);
    m_port =
        atoi(host_data.substr(m_idx + 1, host_data.size() - m_idx).c_str());
    SrpcLogger::INFO("service name: " + m_service_name + " ip: " + m_ip +
                     " port: " + std::to_string(m_port));
    // 尝试连接服务器
    auto rt = newConnect(m_ip.c_str(), m_port);
    if (!rt) {
      SrpcLogger::ERROR("connect server error");
      return;
    } else {
      SrpcLogger::INFO("connect server success");
    }
  }

  // 请求参数序列化
  uint32_t args_size{0};
  std::string args_str;
  if (request->SerializeToString(&args_str)) {
    args_size = args_str.size();
  } else {
    controller->SetFailed("serialize request fail");
    return;
  }

  Srpc::RpcHeader SrpcHeader;
  SrpcHeader.set_service_name(m_service_name);
  SrpcHeader.set_method_name(m_method_name);
  SrpcHeader.set_args_size(args_size);

  // 将RPC头部信息序列化为字符串
  uint32_t header_size{0};
  std::string rpc_header_str;
  if (SrpcHeader.SerializeToString(&rpc_header_str)) {
    header_size = rpc_header_str.size();
  } else {
    controller->SetFailed("serialize rpc header fail");
    return;
  }
  // 将头部长度和头部信息拼接成完整的RPC请求报文
  std::string send_rpc_str;

  google::protobuf::io::StringOutputStream string_output(&send_rpc_str);
  google::protobuf::io::CodedOutputStream coded_output(&string_output);
  // 写入头部长度
  coded_output.WriteVarint32(static_cast<uint32_t>(header_size));
  // 写入头部信息
  coded_output.WriteString(rpc_header_str);
  // 拼接请求参数
  coded_output.WriteString(args_str);
  // 发送RPC请求到服务器
  if (send(m_clientfd, send_rpc_str.c_str(), send_rpc_str.size(), 0) == -1) {
    close(m_clientfd);
    char error_text[512] = {0};
    SrpcLogger::ERROR(std::string("send error: ") +
                      strerror_r(errno, error_text, sizeof(error_text)));
    controller->SetFailed(error_text);
    return;
  }
  // 接收服务的响应
  char recv_buf[1024] = {0};
  int recv_size = 0;
  if ((recv_size = recv(m_clientfd, recv_buf, 1024, 0)) == -1) {
    char error_text[512] = {0};
    SrpcLogger::ERROR(std::string("recv error: ") +
                      strerror_r(errno, error_text, sizeof(error_text)));
    controller->SetFailed(error_text);
    return;
  }
  // 将接收到的响应数据反序列化
  if (!response->ParseFromArray(recv_buf, recv_size)) {
    close(m_clientfd);
    char error_text[512] = {0};
    SrpcLogger::ERROR(std::string("parse error: ") +
                      strerror_r(errno, error_text, sizeof(error_text)));
    controller->SetFailed(error_text);
    return;
  }
  close(m_clientfd);
}

bool SrpcChannel::newConnect(const std::string &ip, uint16_t port) {
  // 创建socket
  int clientFd = socket(AF_INET, SOCK_STREAM, 0);
  if (clientFd == -1) {
    char error_text[512] = {0};
    SrpcLogger::ERROR(std::string("socket error: ") +
                      strerror_r(errno, error_text, sizeof(error_text)));
    return false;
  }
  // 设置服务器地址信息
  struct sockaddr_in server_addr;
  server_addr.sin_family = AF_INET;
  server_addr.sin_port = htons(port);
  server_addr.sin_addr.s_addr = inet_addr(ip.c_str());
  // 尝试连接服务器
  if (connect(clientFd, (struct sockaddr *)&server_addr, sizeof(server_addr)) ==
      -1) {
    close(clientFd);
    char error_text[512] = {0};
    SrpcLogger::ERROR(std::string("connect server error: ") +
                      strerror_r(errno, error_text, sizeof(error_text)));
    return false;
  }
  m_clientfd = clientFd;
  return true;
}

std::string SrpcChannel::QueryServiceHost(ZooKeeperClient *zkclient,
                                          const std::string &service_name,
                                          const std::string &method_name,
                                          int &idx) {
  // 查询路径
  std::string method_path = "/" + service_name + "/" + method_name;
  SrpcLogger::INFO("method path: " + method_path);
  std::unique_lock<std::mutex> lock(g_data_mutex);
  std::string host_data = zkclient->GetData(method_path.c_str());
  lock.unlock();
  if (host_data.empty()) {
    SrpcLogger::INFO(method_path + " is not exist!");
    return "";
  }
  idx = host_data.find(":");
  if (idx == -1) {
    SrpcLogger::INFO(method_path + "address is invalid!");
    return "";
  }
  return host_data;
}